package com.crazyteam.datasource.dynamic;

import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.keygen.SelectKeyGenerator;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author guoyang
 * @Date 2018-01-09
 * @Desc 数据源支持注解自动装配
 */
@Intercepts({
        @Signature(type = Executor.class, method = "update", args = {
                MappedStatement.class, Object.class }),
        @Signature(type = Executor.class, method = "query", args = {
                MappedStatement.class, Object.class, RowBounds.class,
                ResultHandler.class }) })
public class DynamicDateSourceMybatisPlugin implements Interceptor {

    protected static final Logger logger = LoggerFactory.getLogger(DynamicDateSourceMybatisPlugin.class);

    private static final String REGEX = ".*insert\\u0020.*|.*delete\\u0020.*|.*update\\u0020.*";

    private static String selectDefualtDataSource = DynamicDataSourceHolder.SLAVE;

    private static final Map<String, String> cacheMap = new ConcurrentHashMap<>();

    private List<String> slaveSqls = new ArrayList<>();

    private List<String> masterSqls = new ArrayList<>();

    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        String dataSourceName = DynamicDataSourceHolder.getDataSource();


        if (!DynamicDataSourceHolder.MASTER.equalsIgnoreCase(dataSourceName)){

            boolean synchronizationActive = TransactionSynchronizationManager.isSynchronizationActive();


            if(!synchronizationActive) {
                Object[] objects = invocation.getArgs();
                MappedStatement ms = (MappedStatement) objects[0];
                dataSourceName = cacheMap.get(ms.getId());
                if(dataSourceName == null) { // TODO

                    //读方法
                    if(ms.getSqlCommandType().equals(SqlCommandType.SELECT)) {
                        //!selectKey 为自增id查询主键(SELECT LAST_INSERT_ID() )方法，使用主库
                        if(ms.getId().contains(SelectKeyGenerator.SELECT_KEY_SUFFIX)) {
                            dataSourceName = DynamicDataSourceHolder.MASTER;
                        } else {
                            BoundSql boundSql = ms.getSqlSource().getBoundSql(objects[1]);
                            String sql = boundSql.getSql().toLowerCase(Locale.CHINA).replaceAll("[\\t\\n\\r]", " ");
                            if(sql.matches(REGEX)) {
                                dataSourceName = DynamicDataSourceHolder.MASTER;
                            } else {
                                if (!CollectionUtils.isEmpty(masterSqls) && masterSqls.contains(ms.getId())){
                                    dataSourceName = DynamicDataSourceHolder.MASTER;
                                }else if (!CollectionUtils.isEmpty(slaveSqls) && slaveSqls.contains(ms.getId())){
                                    dataSourceName = DynamicDataSourceHolder.SLAVE;
                                }else{
                                    dataSourceName = selectDefualtDataSource;
                                }
                            }
                        }
                    }else{
                        dataSourceName = DynamicDataSourceHolder.MASTER;
                    }
                    logger.debug(String.format("mapper method:[{%s}] use [{%s}] Strategy, SqlCommandType [{%s}]..", ms.getId(), dataSourceName, ms.getSqlCommandType().name()));
                    cacheMap.put(ms.getId(), dataSourceName);
                }
//                else{
//                    dataSourceName = DynamicDataSourceHolder.MASTER;
//                }
            }else{
                dataSourceName = DynamicDataSourceHolder.MASTER;
            }
            DynamicDataSourceHolder.putDataSource(dataSourceName);
        }
        try {
            return invocation.proceed();
        }finally {
            try {
                boolean synchronizationActive = TransactionSynchronizationManager.isSynchronizationActive();
                if (!synchronizationActive){
                    DynamicDataSourceHolder.clearDataSource();
                }
            }catch (Exception e){
            }

        }
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof Executor) {
            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

    @Override
    public void setProperties(Properties properties) {
        //
        if (properties.keySet().contains("selectDefaultDataSource")) {
            if ("MASTER".equalsIgnoreCase(properties.get("selectDefaultDataSource").toString())
                    ||"MASTER".equalsIgnoreCase(properties.get("selectDefaultDataSource").toString())){
                selectDefualtDataSource = DynamicDataSourceHolder.MASTER;
            }
        }
        if (properties.keySet().contains("slaveSqls")) {
            slaveSqls = (List)properties.get("slaveSqls");
        }
        if (properties.keySet().contains("masterSqls")) {
            masterSqls = (List)properties.get("masterSqls");
        }
    }

}
